import sys
import traceback
from typing import Callable, Optional, TypeVar, Union, overload

from zkl_registries.basics import AnyTagKey, AnyTagVal, MutableRegistry, Queries, Tags

AnyCallable = TypeVar('AnyCallable', bound=Callable)


class CallablesMatched(tuple[AnyCallable, ...]):
    def __call__(self, *args, **kwargs):
        impls_error = []
        for impl in reversed(self):
            try:
                return impl(*args, **kwargs)
            except NotImplementedError as err:
                impls_error.append((impl, err))

        if not impls_error:
            raise NotImplementedError(f"No matched impl.")

        if len(impls_error) == 1:
            raise impls_error[0]

        print(f"Matched {len(impls_error)} impls, but all of them raise NotImplementedError.", file=sys.stderr)
        for impl, impl_error in impls_error:
            print(f"\nError from {impl}:", file=sys.stderr)
            traceback.print_exception(impl_error)
        raise NotImplementedError(f"Failed to call all impls.")


class CallablesRegistry(MutableRegistry[AnyTagKey, AnyTagVal, AnyCallable]):
    def __init__(self, base: MutableRegistry[AnyTagKey, AnyTagVal, AnyCallable]):
        self.base = base

    # noinspection PyMethodOverriding
    @overload
    def register(self, tags: Tags) -> Callable[[AnyCallable], AnyCallable]:
        pass

    @overload
    def register(self, tags: Tags, value: AnyCallable) -> AnyCallable:
        pass

    def register(self,
        tags: Tags,
        value: Optional[AnyCallable] = None,
    ) -> Union[Callable[[AnyCallable], AnyCallable], AnyCallable]:
        if value is None:
            def decorated_func(value: AnyCallable) -> AnyCallable:
                self.register(tags, value)
                return value

            return decorated_func
        else:
            self.base.register(tags, value)
            return value

    def match(self, queries: Queries) -> Union[AnyCallable, tuple[AnyCallable, ...]]:
        return CallablesMatched(self.base.match(queries))
